#ifndef _REGULAR_TREE_H_
#define _REGULAR_TREE_H_

#ifdef _DEBUG
#include "Assert.h"
#endif

#include <string.h>

#include "cuTypes.h"
#include "cuMath.h"

#include "Box.h"

template <int nSubNode>
struct TNode
{
    int m_Start;
    int m_Count;

    TNode * m_ppNode[nSubNode];
};

template <int nSubNode>
static TNode<nSubNode> * Node(TNode<nSubNode> * pNode)
{
    for (int iNode = 0; iNode < nSubNode; iNode++)
        pNode->m_ppNode[iNode] = NULL;

    pNode->m_Count = 0;
    pNode->m_Start = 0;

    return pNode;
}

template <int nSubNode>
static void Release(TNode<nSubNode> * pNode)
{
    if (pNode != NULL)
    {
        if (pNode->m_ppNode != NULL)
        {
            for (int iNode = 0; iNode < nSubNode; iNode++)
            {
                Release(pNode->m_ppNode[iNode]);

                delete pNode->m_ppNode[iNode];
            }
        }
    }
};

template <class T, int DimX, int DimY, int DimZ>
struct THierarchicalTree
{   
    S_Box m_AABB;

    uint32 m_MaxDepth;
    uint32 m_MaxElem;

    T * m_pData;

    TNode<DimX * DimY * DimZ> * m_pRoot;
};

template <class T, int DimX, int DimY, int DimZ>
static THierarchicalTree<T, DimX, DimY, DimZ> * HierarchicalTree(THierarchicalTree<T, DimX, DimY, DimZ> *pTree, int maxDepth = 1)
{
    if (pTree != NULL)
    {
        pTree->m_MaxDepth = maxDepth;

        AABox(&pTree->m_AABB);   

        pTree->m_pRoot = Node(new TNode<DimX * DimY * DimZ>);

        pTree->m_pData = NULL;
    }

    return pTree;
}

template <class T, int DimX, int DimY, int DimZ>
static void Release(THierarchicalTree<T, DimX, DimY, DimZ> *pTree)
{
    if (pTree != NULL)
    {
        if (pTree->m_pRoot != NULL)
        {
            Release(pTree->m_pRoot);

            delete pTree->m_pRoot;
        }

        if (pTree->m_pData != NULL)
            delete [] pTree->m_pData;
    }
}

template <int nSubNode>
static void Scan(TNode<nSubNode> *pNode, int & scanIdx, int depth, int nMaxDepth)
{
    if (pNode != NULL)
    {
        if (depth <= nMaxDepth)
        {
            pNode->m_Start = scanIdx;

            scanIdx += pNode->m_Count;

            pNode->m_Count = 0;

            for (int iSubNode = 0; iSubNode < nSubNode; iSubNode++)
                Scan(pNode->m_ppNode[iSubNode], scanIdx, depth+1, nMaxDepth);
        }
    }
}

// ----------------------------------------------------------------------------
// * General
// THierarchicalTree operates in three stages that are encapsulated in Build()
// Build receives a pointer pElem to objects you want to organize in a tree
// It also receives nElem which is number of elements stored in pElem
// Build has 3 Stages:
// * Stages
// ** Count():
// *** For each element Count() checks which Node it belongs to
// *** Increments the m_Count value for this node
// *** An element belongs to a Node if the following is true:
// **** all points for the element belong to Nodes aabb                     AND
// **** this is the last possible depth level (limited by maxDepth)         OR
// **** an element fits into this level but doesn't fit into the next level
// ** Scan():
// *** Nodes are not dynamic structures
// *** Tree holds a pointer to tree elements
// *** Tree allocates and releases memory for this pointer
// *** Each Node wants to take a piece of this memory for it's elements
// *** So Scan runs through each node and calculates Start index for each Node
// ** Insert():
// *** Insert() is almost a copy of Count()
// *** For each element Insert() does all the math as in Count()
// *** When the host Node is found - element is inserted at address: 
// *** Tree->m_pData[pNode->Start + pNode->Count]
// *** Which is a correct offset to trees m_pData array of elements
// 
// In order for an element to be insertable into a tree it needs to implement a
// CompareTo() function that evaluates which aabb the elemnt falls into
// ----------------------------------------------------------------------------


// ----------------------------------------------------------------------------
// Implementation for float3 values
// ----------------------------------------------------------------------------

template <int nSubNode>
static void Insert(TNode<nSubNode> * pNode, float3 * pData, S_Box aabb, int3 dim, float3 point, int depth, int maxDepth)
{
    if (pNode != NULL)
    {
        if (depth+1 <= maxDepth)
        {         
            int3 nodeId = CompareTo(aabb, point, dim);

            int NodeIdx = nodeId.x + (nodeId.y + nodeId.z * dim.y) * dim.x;

            if (NodeIdx >= 0 && NodeIdx < nSubNode)
            {               
                if (pNode->m_ppNode[NodeIdx] == NULL)
                    pNode->m_ppNode[NodeIdx] = Node(new TNode<nSubNode>);

                float3 l3 = aabb.pmax - aabb.pmin;

                float3 invDim = { 1.0f / dim.x, 1.0f / dim.y, 1.0f / dim.z};

                float3 pmin = aabb.pmin;
                pmin.x += nodeId.x * l3.x * invDim.x;
                pmin.y += nodeId.y * l3.y * invDim.y;
                pmin.z += nodeId.z * l3.z * invDim.z;

                float3 pmax = aabb.pmax;
                pmax.x -= (dim.x - 1 - nodeId.x) * l3.x * invDim.x;
                pmax.y -= (dim.y - 1 - nodeId.y) * l3.y * invDim.y;
                pmax.z -= (dim.z - 1 - nodeId.z) * l3.z * invDim.z;

                AABox(&aabb, pmin, pmax);

                Insert(pNode->m_ppNode[NodeIdx], 
                       pData,
                       aabb, 
                       dim,
                       point, 
                       depth+1,
                       maxDepth);
            }
            else
            {
                pData[pNode->m_Start + pNode->m_Count] = point;
                pNode->m_Count++;
            }
                
        }
        else
        {
            pData[pNode->m_Start + pNode->m_Count] = point;
            pNode->m_Count++;
        }
    }
}

template <int nSubNode>
static void Zero(TNode<nSubNode> * pNode)
{
    if (pNode != NULL)
    {         
        pNode->m_Count = 0;
        pNode->m_Start = 0;

        for (int iNode = 0; iNode < nSubNode; iNode++)
        {
            if (pNode->m_ppNode[iNode] != NULL)
                Zero(pNode->m_ppNode[iNode]);
        }
    }
    
}

template <int nSubNode>
static void Count(TNode<nSubNode> * pNode, S_Box aabb, int3 dim, float3 point, int depth, int maxDepth)
{
    if (pNode != NULL)
    {
        if (depth+1 <= maxDepth)
        {         
            int3 nodeId = CompareTo(aabb, point, dim);

            int NodeIdx = nodeId.x + (nodeId.y + nodeId.z * dim.y) * dim.x;

            if (NodeIdx >= 0 && NodeIdx < nSubNode)
            {               
                if (pNode->m_ppNode[NodeIdx] == NULL)
                    pNode->m_ppNode[NodeIdx] = Node(new TNode<nSubNode>);

                float3 l3 = aabb.pmax - aabb.pmin;

                float3 invDim = { 1.0f / dim.x, 1.0f / dim.y, 1.0f / dim.z};

                float3 pmin = aabb.pmin;
                pmin.x += nodeId.x * l3.x * invDim.x;
                pmin.y += nodeId.y * l3.y * invDim.y;
                pmin.z += nodeId.z * l3.z * invDim.z;

                float3 pmax = aabb.pmax;
                pmax.x -= (dim.x - 1 - nodeId.x) * l3.x * invDim.x;
                pmax.y -= (dim.y - 1 - nodeId.y) * l3.y * invDim.y;
                pmax.z -= (dim.z - 1 - nodeId.z) * l3.z * invDim.z;

                AABox(&aabb, pmin, pmax);

                Count(pNode->m_ppNode[NodeIdx], aabb, dim, point, depth+1, maxDepth);
            }
            else
                pNode->m_Count++;
        }
        else
            pNode->m_Count++;
    }
}

template <int DimX, int DimY, int DimZ>
static void Build(THierarchicalTree<float3, DimX, DimY, DimZ> *pTree, float3 * pElem, uint32 nElem)
{
    if (pTree != NULL)
    {
        Zero(pTree->m_pRoot);

        int3 Dim = {DimX, DimY, DimZ};
        for (uint32 iElem = 0; iElem < nElem; iElem++)
        {
            Count(pTree->m_pRoot,
                  pTree->m_AABB,
                  Dim,
                  pElem[iElem],
                  0,
                  pTree->m_MaxDepth);
        }

        int scanIdx = 0;
        int depthIdx = 0;

        Scan(pTree->m_pRoot, scanIdx, depthIdx, pTree->m_MaxDepth);

        if (pTree->m_pData != NULL && pTree->m_MaxElem < nElem)
        {
            delete [] pTree->m_pData;
            pTree->m_pData = NULL;
            pTree->m_MaxElem = 0;
        }       

        if (pTree->m_pData == NULL)
        {
            pTree->m_pData = new float3[nElem];
            memset(pTree->m_pData, 0, sizeof(float3) * nElem);
        }

        pTree->m_MaxElem = nElem;

        for (uint32 iElem = 0; iElem < nElem; iElem++)
        {
            Insert(pTree->m_pRoot, 
                   pTree->m_pData, 
                   pTree->m_AABB,
                   Dim,
                   pElem[iElem],
                   0,
                   pTree->m_MaxDepth);
        }
    }
}

// ----------------------------------------------------------------------------
// Implementation for CXobjIO values
// ----------------------------------------------------------------------------

template <int nSubNode>
static void Insert(TNode<nSubNode> * pNode, int4 * pData, S_Box aabb, int3 dim, float3 v[3], int4 idx, int depth, int maxDepth)
{
    if (pNode != NULL)
    {
        if (depth+1 <= maxDepth)
        {         
            int3 nodeId = CompareTo(aabb, v, dim);

            int NodeIdx = nodeId.x + (nodeId.y + nodeId.z * dim.y) * dim.x;

            if (NodeIdx >= 0 && NodeIdx < nSubNode)
            {               
                if (pNode->m_ppNode[NodeIdx] == NULL)
                    pNode->m_ppNode[NodeIdx] = Node(new TNode<nSubNode>);

                float3 l3 = aabb.pmax - aabb.pmin;

                float3 invDim = { 1.0f / dim.x, 1.0f / dim.y, 1.0f / dim.z};

                float3 pmin = aabb.pmin;
                pmin.x += nodeId.x * l3.x * invDim.x;
                pmin.y += nodeId.y * l3.y * invDim.y;
                pmin.z += nodeId.z * l3.z * invDim.z;

                float3 pmax = aabb.pmax;
                pmax.x -= (dim.x - 1 - nodeId.x) * l3.x * invDim.x;
                pmax.y -= (dim.y - 1 - nodeId.y) * l3.y * invDim.y;
                pmax.z -= (dim.z - 1 - nodeId.z) * l3.z * invDim.z;

                AABox(&aabb, pmin, pmax);

                Insert(pNode->m_ppNode[NodeIdx], 
                       pData, aabb, dim, v,
                       idx, depth+1, maxDepth);
            }
            else
            {
                pData[pNode->m_Start + pNode->m_Count] = idx;
                pNode->m_Count++;
            }

        }
        else
        {
            pData[pNode->m_Start + pNode->m_Count] = idx;
            pNode->m_Count++;
        }
    }
}

template <int nSubNode>
static void Count(TNode<nSubNode> * pNode, S_Box aabb, int3 dim, float3 v[3], int depth, int maxDepth)
{
    if (pNode != NULL)
    {
        if (depth+1 <= maxDepth)
        {         
            int3 nodeId = CompareTo(aabb, v, dim);

            int NodeIdx = nodeId.x + (nodeId.y + nodeId.z * dim.y) * dim.x;

            if (NodeIdx >= 0 && NodeIdx < nSubNode)
            {               
                if (pNode->m_ppNode[NodeIdx] == NULL)
                    pNode->m_ppNode[NodeIdx] = Node(new TNode<nSubNode>);

                float3 l3 = aabb.pmax - aabb.pmin;

                float3 invDim = { 1.0f / dim.x, 1.0f / dim.y, 1.0f / dim.z};

                float3 pmin = aabb.pmin;
                pmin.x += nodeId.x * l3.x * invDim.x;
                pmin.y += nodeId.y * l3.y * invDim.y;
                pmin.z += nodeId.z * l3.z * invDim.z;

                float3 pmax = aabb.pmax;
                pmax.x -= (dim.x - 1 - nodeId.x) * l3.x * invDim.x;
                pmax.y -= (dim.y - 1 - nodeId.y) * l3.y * invDim.y;
                pmax.z -= (dim.z - 1 - nodeId.z) * l3.z * invDim.z;

                AABox(&aabb, pmin, pmax);

                Count(pNode->m_ppNode[NodeIdx], aabb, dim, v, depth+1, maxDepth);
            }
            else
                pNode->m_Count++;
        }
        else
            pNode->m_Count++;
    }
}

template <int DimX, int DimY, int DimZ>
static void Build(THierarchicalTree<int4, DimX, DimY, DimZ> *pTree, C_Actor * pXobj, uint32 nXobj)
{
    if (pTree != NULL)
    {
        Zero(pTree->m_pRoot);

        int3 Dim = {DimX, DimY, DimZ};

        uint32 nElem = 0;
        for (uint32 iXobj = 0; iXobj < nXobj; iXobj++)
            for (uint32 iFace = 0; iFace < pXobj[iXobj].posIdx.size(); iFace+=3)
            {
                float3 v[3];
                int4   idx = { pXobj[iXobj].posIdx[iFace + 0],
                               pXobj[iXobj].posIdx[iFace + 1], 
                               pXobj[iXobj].posIdx[iFace + 2],
                               iXobj };

                if (idx.x == idx.y || idx.x == idx.z || idx.y == idx.z)
                    continue;

                v[0] = xyz( pXobj[iXobj].pos[idx.x] );
                v[1] = xyz( pXobj[iXobj].pos[idx.y] );
                v[2] = xyz( pXobj[iXobj].pos[idx.z] );

                Count(pTree->m_pRoot,
                      pTree->m_AABB,
                      Dim, v, 0,
                      pTree->m_MaxDepth);

                nElem++;
            }

        int scanIdx = 0;
        int depthIdx = 0;

        Scan(pTree->m_pRoot, scanIdx, depthIdx, pTree->m_MaxDepth);

        if (pTree->m_pData != NULL && pTree->m_MaxElem < nElem)
        {
            delete [] pTree->m_pData;
            pTree->m_pData = NULL;
            pTree->m_MaxElem = 0;
        }       

        if (pTree->m_pData == NULL)
        {
            pTree->m_pData = new int4[nElem];
            memset(pTree->m_pData, 0, sizeof(int4) * nElem);
        }

        pTree->m_MaxElem = nElem;

        for (uint32 iXobj = 0; iXobj < nXobj; iXobj++)
            for (uint32 iFace = 0; iFace < pXobj[iXobj].posIdx.size(); iFace+=3)
            {
                float3 v[3];
                int4   idx = {pXobj[iXobj].posIdx[iFace],
                              pXobj[iXobj].posIdx[iFace+1], 
                              pXobj[iXobj].posIdx[iFace+2],
                              iXobj};

                if (idx.x == idx.y || idx.x == idx.z || idx.y == idx.z)
                    continue;

                v[0] = xyz( pXobj[iXobj].pos[idx.x] );
                v[1] = xyz( pXobj[iXobj].pos[idx.y] );
                v[2] = xyz( pXobj[iXobj].pos[idx.z] );

                Insert(pTree->m_pRoot, 
                    pTree->m_pData, 
                    pTree->m_AABB,
                    Dim, v, idx, 0,
                    pTree->m_MaxDepth);
            }
    }
}

typedef THierarchicalTree<float3, 2, 2, 2>    float3Octree;
typedef THierarchicalTree<int4,   2, 2, 2>    int4Octree;


#endif